#!/usr/bin/env python3


import glob
import pathlib
import os
import json
import random
import sys
import SimpleITK as sitk
import numpy as np
import nibabel as nib
import shutil
import pandas as pd
from scipy.ndimage import morphology
from collections import Counter
import cv2
import math
import yaml


# def _get_pairs(pathstr, img_pattern, mask_patterns):
#     """Get image-label pair from sepcific label_patterns and a single image_pattern.

#     Args:
#         pathstr (str): dataset path. e.g, '/home/yusongli/_dataset/_IIPL/ShuaiWang/20211214/shidaoai/sichuan'
#         img_pattern (str): e.g, '*CT*.gz'
#         label_patterns (list[str]): label patterns. e.g, label_patterns=['*T[_1 ]*.gz']

#     Returns:
#         (image, label): yield
#     """
#     path = pathlib.Path(pathstr)
#     for pattern in mask_patterns:
#         for item in list(path.rglob(pattern)):
#             if not item.exists():
#                 continue
#             label = item.as_posix()
#             image = list(item.parent.rglob(img_pattern))
#             if not image:
#                 continue
#             image = image[0].as_posix()
#             yield image, label




def no_label():
    from monai.transforms import (
        AddChanneld,
        Compose,
        LoadImaged,
        ToTensord,
    )

    from monai.data import (
        DataLoader,
        CacheDataset,
        load_decathlon_datalist,
    )

    import pathlib

    orig_transforms = Compose(
        [
            LoadImaged(keys=['image', 'label']),
            AddChanneld(keys=['image', 'label']),
            ToTensord(keys=['image', 'label'])
        ]
    )

    split_JSON = "dataset.json"
    datasets = split_JSON

    train_files = load_decathlon_datalist(datasets, True, "all")

    orig_ds = CacheDataset(
        data=train_files,
        transform=orig_transforms,
        cache_num=24,
        cache_rate=1.0,
        num_workers=8,
    )

    # 未做变换的训练集
    orig_loader = DataLoader(orig_ds, batch_size=1, shuffle=False, num_workers=8, pin_memory=True)

    # no_label = [item['image_meta_dict']['filename_or_obj'][0] for item in orig_loader if item['label'].sum() == 0 ]

    with open(pathlib.Path(log_path).as_posix(), 'w') as f:
        for item in orig_loader:
            if item['label'].sum() == 0:
                print(item['label_meta_dict']['filename_or_obj'][0])
                f.write(item['label_meta_dict']['filename_or_obj'][0] + '\n')
                f.flush()

    # with open(pathlib.Path('dataset', 'no_label.txt').as_posix(), 'w') as f:
    #     for item in no_label:
    #         f.write(item + '\n')


def _get_roi(img_path=None, mask_path=None, pred_path=None, img_savepath=None, mask_savepath=None, x_area=None, y_area=None, z_area=None, scale=False, zoom=None, translation=True, mode='even', debug_embed_mode=0, console=True, save=False, log_path=None):
    """Get ROI from a specific image and mask.
    Args:
        img_path (str): The path of the image. If None, the images will be ignored and only masks will be processed.
        mask_path (str): The path of the mask.
        pred_path (str, default None): If not None, cut the img_path and mask_path according to the pred_path.
        img_savepath (str): The path of the image after ROI.
        mask_savepath (str): The path of the mask after ROI.
        x_area (int): The x-axis area of ROI. If None, the tumor contour will be used.
        y_area (int): The y-axis area of ROI. If None, the tumor contour will be used.
        z_area (int): The z-axis area of ROI. If None, the tumor contour will be used.
        scale (bool): If True, the img and mask will be scaled intensity.
        zoom (float): if not None, the ROI radius will be zoomed (directly multiple) by this value. So if this value larger than 0, the ROI radius will be larger. Otherwise smaller.
        console (bool): If True, print the path into console.
        translation (bool): Judje if translate when bounding box is out of image margin.
        mode (str): 'edge', 'even', 'max' for different cutting strategy.
        debug_embed_mode (int): different embedding rate caculate stretegy.
        output (bool, default=True): print current processing image path to console.
        save (bool, default=False): If True, save offline, else not save (online).
        log_path (str): log file path.
    Return:
        x_radius (int): The x-axis radius of ROI.
        y_radius (int): The y-axis radius of ROI.
        z_radius (int): The z-axis radius of ROI.
    """

    # Get save path.
    if img_savepath is not None and mask_savepath is not None:
        img_savepath = pathlib.Path(img_savepath)
        mask_savepath = pathlib.Path(mask_savepath)

    if console:
        print(f'{img_path} | {mask_path}')

    # Load img and mask.
    mask_path = pathlib.Path(mask_path)

    have_img_flag = True
    try:
        img_path = pathlib.Path(img_path)
        img = nib.load(img_path.as_posix())
        img_array = img.get_fdata().copy()
        img_shape = img_array.shape
    except:
        have_img_flag = False

    mask = nib.load(mask_path.as_posix())
    if debug_embed_mode != 0:
        mask_array_bak = mask.get_fdata().copy()
    mask_array = mask.get_fdata().copy()
    mask_shape = mask_array.shape
    target_array = mask_array

    if have_img_flag:
        for i in range(len(img_shape)):
            if img_shape[i] != mask_shape[i]:
                # Img does not match with mask. They are in different shape.
                return -1

    if pred_path:
        pred_path = pathlib.Path(pred_path)
        pred = nib.load(pred_path.as_posix())
        pred_array = pred.get_fdata().copy()
        pred_shape = pred_array.shape
        for i in range(len(mask_shape)):
            if mask_shape[i] != pred_shape[i]:
                # Mask does not match with pred. They are in different shape.
                return -2
        target_array = pred_array

    nozero = np.nonzero(target_array)

    try:
        # The first dimension: x
        # The second dimension: y
        # The third dimension: z
        x_min = nozero[0].min()
        x_max = nozero[0].max()
        y_min = nozero[1].min()
        y_max = nozero[1].max()
        z_min = nozero[2].min()
        z_max = nozero[2].max()
    except:
        # No mask.
        return -3

    x_width = x_max - x_min + 1
    y_width = y_max - y_min + 1
    z_width = z_max - z_min + 1

    x_radius = math.ceil(x_width / 2)
    y_radius = math.ceil(y_width / 2)
    z_radius = math.ceil(z_width / 2)

    x_center = x_min + x_radius
    y_center = y_min + y_radius
    z_center = z_min + z_radius

    # If the size of the roi is artificially specified, the specified value is used. Otherwise, cut along the tumor boundary.
    # 1. Check if x_area is larger than the radius of the image. If larger, the x_radius will be set to the image radius

    x_radius = x_area if x_area else x_radius
    y_radius = y_area if y_area else y_radius
    z_radius = z_area if z_area else z_radius

    if zoom is not None:
        x_radius = math.ceil(x_radius * zoom)
        y_radius = math.ceil(y_radius * zoom)
        z_radius = math.ceil(z_radius * zoom)

    max_radius = max(x_radius, y_radius, z_radius)

    # ! <<< Different cutting method.
    # Smallest ROI, cut along edge in different axis.
    if mode == 'edge':
        x_roi_min_orig = x_min
        x_roi_max_orig = x_max
        y_roi_min_orig = y_min
        y_roi_max_orig = y_max
        z_roi_min_orig = z_min
        z_roi_max_orig = z_max
    # Middle roi, but left radius and right radius are same (keep larger one).
    elif mode == 'even':
        x_roi_min_orig = x_center - x_radius
        x_roi_max_orig = x_center + x_radius
        y_roi_min_orig = y_center - y_radius
        y_roi_max_orig = y_center + y_radius
        z_roi_min_orig = z_center - z_radius
        z_roi_max_orig = z_center + z_radius
    # Common max roi
    elif mode == 'max':
        x_roi_min_orig = x_center - (max_radius + 1)
        x_roi_max_orig = x_center + (max_radius + 1)
        y_roi_min_orig = y_center - (max_radius + 1)
        y_roi_max_orig = y_center + (max_radius + 1)
        z_roi_min_orig = z_center - (max_radius + 1)
        z_roi_max_orig = z_center + (max_radius + 1)
    # ! >>>

    # Note: If the voxal concatinates with the image boundary, there will be a bug that the cropped array will not be a cube. It need to be fixed in the future.
    # ? <<< Fix bug
    if not translation:
        x_roi_min_norm = max(x_roi_min_orig, 0)
        x_roi_max_norm = min(x_roi_max_orig, mask_shape[0])
        y_roi_min_norm = max(y_roi_min_orig, 0)
        y_roi_max_norm = min(y_roi_max_orig, mask_shape[1])
        z_roi_min_norm = max(z_roi_min_orig, 0)
        z_roi_max_norm = min(z_roi_max_orig, mask_shape[2])
    else:
        def fix_offset(roi_min_orig, roi_max_orig, img_shape):
            if roi_max_orig > (img_shape - 1):
                roi_min_orig -= abs(roi_max_orig - (img_shape - 1))
                roi_max_orig = img_shape - 1
            if roi_min_orig < 0:
                roi_max_orig += abs(roi_min_orig)
                roi_min_orig = 0
            # print(roi_min_orig, roi_max_orig, (roi_max_orig - roi_min_orig))
            return roi_min_orig, roi_max_orig

        x_roi_min_norm, x_roi_max_norm = fix_offset(x_roi_min_orig, x_roi_max_orig, mask_array.shape[0])
        y_roi_min_norm, y_roi_max_norm = fix_offset(y_roi_min_orig, y_roi_max_orig, mask_array.shape[1])
        z_roi_min_norm, z_roi_max_norm = fix_offset(z_roi_min_orig, z_roi_max_orig, mask_array.shape[2])
    # ? >>>

    mask_roi_array = mask_array[x_roi_min_norm:x_roi_max_norm, y_roi_min_norm:y_roi_max_norm, z_roi_min_norm:z_roi_max_norm]

    if debug_embed_mode == 1:
        big_array = np.zeros_like(mask_array_bak)
        small_array = mask_roi_array
        big_index = (x_roi_min_orig, y_roi_min_orig, z_roi_min_orig)
        big_array = embed(small_array, big_array, big_index)
        dice_value = dice(big_array, mask_array_bak)
        return dice_value
    elif debug_embed_mode == 2:
        small_voxel = np.sum(mask_roi_array)
        big_voxel = np.sum(mask_array_bak)
        return small_voxel / big_voxel

    # Debug:
    # print(mask_roi_array.shape)
    # if mask_roi_array.shape[0] != 224 or mask_roi_array.shape[1] != 224 or mask_roi_array.shape[2] <= 0 or mask_roi_array.shape[2] > 96:
    #     print(f'Error occured: {img_path}, orig shape: {img_array.shape}')
    #     sys.exit()

    if have_img_flag:
        img_roi_array = img_array[x_roi_min_norm:x_roi_max_norm, y_roi_min_norm:y_roi_max_norm, z_roi_min_norm:z_roi_max_norm]

        if scale:
            # Image intensity limitation and normalization.
            img_roi_array = _scale_intensity(img_roi_array)

    # if mask_roi_array.shape[0] != mask_roi_array.shape[1] or mask_roi_array.shape[0] != mask_roi_array.shape[2]:
    #     # The cropped area is not a cube. Error generated, return.
    #     return -2

    if img_savepath is not None and mask_savepath is not None and save:
        # ! <<< Save img and mask.
        img_savefolder = img_savepath.parent
        mask_savefolder = mask_savepath.parent

        if have_img_flag and not img_savefolder.exists():
            img_savefolder.mkdir(parents=True, exist_ok=True)
        if not mask_savefolder.exists():
            mask_savefolder.mkdir(parents=True, exist_ok=True)

        if have_img_flag:
            # ? <<<
            # img_out = nib.Nifti1Image(img_roi_array, header= affine=np.eye(4))
            img_out = nib.Nifti1Image(img_roi_array, header=img.header, affine=img.affine)
            # Get image header for fix resolution.
            # img_out.header.set_zooms(img.header.get_zooms())
            img_out.to_filename(img_savepath.as_posix())
            # ? >>>

        # ? <<<
        # mask_out = nib.Nifti1Image(mask_roi_array, affine=np.eye(4))
        mask_out = nib.Nifti1Image(mask_roi_array, header=mask.header, affine=mask.affine)
        # Get mask header for fix resolution.
        # mask_out.header.set_zooms(mask.header.get_zooms())
        mask_out.to_filename(mask_savepath.as_posix())
        # ? >>>
        # ! >>>

    # return max_radius
    return x_radius, y_radius, z_radius


# def get_roi(convert_json_path, tags, keys, input_index, output_index, x_area=None, y_area=None, z_area=None, pred_path=None, scale=False, output=True, save=False, log_path=None):
#     open(log_path , 'w').close()
#     for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
#         img_path = orig_paths[0]
#         mask_path = orig_paths[1]
#         img_savepath = target_paths[0]
#         mask_savepath = target_paths[1]
#         _get_roi(img_path, mask_path, img_savepath=img_savepath, mask_savepath=mask_savepath, x_area=x_area, y_area=y_area, z_area=z_area, pred_path=pred_path, scale=scale, console=output, save=save, log_path=log_path)


# def get_embed(convert_json_path, tags, keys, input_index, output_index=-1, x_area=None, y_area=None, z_area=None, pred_path=None, scale=False, zoom=None, translation=True, mode='even', debug_embed_mode=0, output=True, save=False, log_path=None):
# # img_path, mask_path, img_savepath=None, mask_savepath=None, x_area=None, y_area=None, z_area=None, pred_path=None, scale=False, zoom=None, translation=True, mode='even', debug_embed=False, output=True, save=False, log_path=None
#     open(log_path , 'w').close()
#     result = []
#     i = 0
#     for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
#         i += 1
#         if i == 5:
#             break
#         img_path = orig_paths[0]
#         mask_path = orig_paths[1]
#         # img_savepath = target_paths[0]
#         # mask_savepath = target_paths[1]
#         result.append(_get_roi(img_path, mask_path, img_savepath=img_savepath, mask_savepath=mask_savepath, x_area=x_area, y_area=y_area, z_area=z_area, pred_path=pred_path, scale=scale, zoom=zoom, translation=translation, mode=mode, debug_embed_mode=debug_embed_mode, output=output, save=save, log_path=log_path))
#     print(sum(result) / len(result))


def scale_intensity(img_path, img_savepath, mask_path=None, mask_savepath=None, norm=True):
    """Get ROI from a specific image.
    Args:
        img_path (str): The path of the image.
        img_savepath (str): The path of the image after ROI.
    Return:
        New image.
    """
    # Get save path.
    img_path = pathlib.Path(img_path)
    img_savepath = pathlib.Path(img_savepath)
    if mask_path is not None and mask_savepath is not None:
        mask_path = pathlib.Path(mask_path)
        mask_savepath = pathlib.Path(mask_savepath)
    try:
        img = nib.load(img_path.as_posix())
        img_array = img.get_fdata().copy()
        print(f'{img_path}')
    except:
        print(f' Not a valid image file: {img_path}')
        return

    # img_array = _scale_intensity(img_array, a_min=0, a_max=1500, b_min=0, b_max=1)
    img_array = _scale_intensity(img_array, a_min=0, a_max=1500, norm=norm)

    # ! <<< Save img.
    img_savefolder = img_savepath.parent

    if not img_savefolder.exists():
        img_savefolder.mkdir(parents=True, exist_ok=True)

    img_out = nib.Nifti1Image(img_array, header=img.header, affine=img.affine)
    # Get image header for fix resolution.
    # img_out.header.set_zooms(img.header.get_zooms())
    img_out.to_filename(img_savepath.as_posix())
    # ! >>>

    # ! <<< Copy mask if need
    if mask_path is not None and mask_savepath is not None:
        if not mask_path.exists():
            return
        mask_savefolder = mask_savepath.parent
        if not mask_savefolder.exists():
            mask_savefolder.mkdir(parents=True, exist_ok=True)
        shutil.copy(mask_path.as_posix(), mask_savepath.as_posix())
    # ! >>>

def _scale_intensity(img_array, a_min=0, a_max=1500, norm=True, b_min=0, b_max=1):
    """Scale the intensity of the image.
    Args:
        img_array: The image array.
        a_min: The minimum value of the original intensity.
        a_max: The maximum value of the original intensity.
        b_min: The minimum value of the new intensity.
        b_max: The maximum value of the new intensity.
    """
    # Image intensity limitation.
    img_array[img_array > a_max] = a_max
    img_array[img_array < a_min] = a_min

    # Image intensity normalization.
    if norm:
        img_array = (img_array - a_min) / (a_max - a_min)
        img_array = img_array * (b_max - b_min) + b_min

    return img_array


def check_pixel_and_dice(data_path, data_path_pattern='**/24/**/', pred_path_pattern='*pred_0*.nii.gz', ct_path_pattern='*CT.nii.gz', mask_path_pattern='*MASK.nii.gz'):
    """Check if the image is no pixel.
    Args:
    Return:
    """
    missing = []
    dice_total = []
    dice_value = {'min': [9999, None], 'max': [-1, None]}
    for pred_path in pathlib.Path(data_path).rglob(data_path_pattern + pred_path_pattern):
        print(pred_path)
        folder = pred_path.parent
        mask_path = folder.glob(mask_path_pattern)

        pred_path = pathlib.Path(pred_path)
        pred = nib.load(pred_path.as_posix())
        pred_array = pred.get_fdata().copy()
        pred_shape = pred_array.shape

        mask_path = pathlib.Path(list(mask_path)[0])
        mask = nib.load(mask_path.as_posix())
        mask_array = mask.get_fdata().copy()
        mask_shape = mask_array.shape

        mydice = dice(pred_array, mask_array)
        dice_total.append(mydice)

        if mydice < dice_value['min'][0]:
            dice_value['min'][0] = mydice
            dice_value['min'][1] = folder.name
        if mydice > dice_value['max'][0]:
            dice_value['max'][0] = mydice
            dice_value['max'][1] = folder.name

        if pred_array.sum() <= 0:
            missing.append(item.as_posix())

    dice_value['avg'] = sum(dice_total) / len(dice_total)

    return missing, dice_value


def check_contrast(datastr, log_path='check_contrast.txt'):
    """Check the contrast of the image.
    Args:
        data_path: The path of the data.
    Return:
        A file record the contrast of the image.
    """
    datapath = pathlib.Path(datastr)
    imgs = datapath.rglob(pattern='*CT*')

    with open(log_path, 'w') as f:
        for img in imgs:
            ct = sitk.ReadImage(img.as_posix())
            ct_array = sitk.GetArrayFromImage(ct)
            print(ct_array.min(), ct_array.max(), img)
            f.write(f'{ct_array.min()}, {ct_array.max()}, {img}\n')
            f.flush()


def _morphology_dilate(orig_img_path, orig_mask_path, target_img_path, target_mask_path, kernel_size=3, iterations=1):
    orig_img_path = pathlib.Path(orig_img_path)
    orig_mask_path = pathlib.Path(orig_mask_path)
    target_img_path = pathlib.Path(target_img_path)
    target_mask_path = pathlib.Path(target_mask_path)

    mask = nib.load(orig_mask_path.as_posix())
    mask_array = mask.get_fdata().copy()

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
    dilated_mask_array = cv2.dilate(mask_array, kernel, iterations=iterations)

    # ! <<< Save dilated mask
    out = nib.Nifti1Image(dilated_mask_array, header=mask.header, affine=mask.affine)
    target_mask_folder = target_mask_path.parent
    if not target_mask_folder.exists():
        target_mask_folder.mkdir(parents=True, exist_ok=True)
    out.to_filename(target_mask_path.as_posix())

    # Copy orig CT img
    shutil.copy(orig_img_path, target_img_path)
    # Copy before dilated mask to save destination
    before_dilated_dst = target_mask_folder.as_posix() + os.sep + orig_mask_path.name.split('.')[0] + '_before_dilated.nii.gz'
    shutil.copy(orig_mask_path, before_dilated_dst)
    # ! >>>


def morphology_dilate(convert_json_path, tags, keys, input_index, output_index, kernel_size=3, iterations=1):
    for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
        orig_img_path = orig_paths[0]
        orig_mask_path = orig_paths[1]
        target_img_path = target_paths[0]
        target_mask_path = target_paths[1]
        _morphology_dilate(orig_img_path, orig_mask_path, target_img_path, target_mask_path, kernel_size=kernel_size, iterations=iterations)


def json_move(convert_json_path, tags, keys, input_index, output_index, mode='copy', cover=False):
    for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
        if len(orig_paths) != len(target_paths):
            print('orig_paths and target_paths should have the same length!')
            return
        for i in range(len(orig_paths)):
            if orig_paths[i] is None or target_paths[i] is None:
                continue
            orig = pathlib.Path(orig_paths[i])
            target = pathlib.Path(target_paths[i])
            if not orig.exists():
                print(f'Orig not exist: {orig}')
                continue
            if cover == False and target.exists():
                print(f'Target already exist: {target}')
                continue
            target_folder = target.parents[0]
            if not target_folder.exists():
                target_folder.mkdir(parents=True, exist_ok=True)
            if mode == 'copy':
                shutil.copy(orig, target)
            elif mode == 'cut':
                shutil.move(orig, target)


def check_affine(convert_json_path, tags, keys, input_index, output_index):
    for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
        img_old_path = orig_paths[0]
        img_new_path = target_paths[0]
        if len(keys) > 1:
            mask_old_path = orig_paths[1]
            mask_new_path = target_paths[1]

        img_old = nib.load(img_old_path)
        # img_old_array = img_old.get_fdata().copy().astype(np.int15)
        # img_old_array = img_old.get_fdata().copy()
        img_old_array = img_old.get_fdata()

        img_new = nib.load(img_new_path)
        # img_new_array = img_new.get_fdata().copy().astype(np.int16)
        # img_new_array = img_new.get_fdata().copy()
        img_new_array = img_new.get_fdata()

        print(np.array_equal(img_old_array, img_new_array))


def dice(im1, im2, empty_score=1.0):
    """
    Computes the Dice coefficient, a measure of set similarity.
    Parameters
    ----------
    im1 : array-like, bool
        Any array of arbitrary size. If not boolean, will be converted.
    im2 : array-like, bool
        Any other array of identical size. If not boolean, will be converted.
    Returns
    -------
    dice : float
        Dice coefficient as a float on range [0,1].
        Maximum similarity = 1
        No similarity = 0
        Both are empty (sum eq to zero) = empty_score

    Notes
    -----
    The order of inputs for `dice` is irrelevant. The result will be
    identical if `im1` and `im2` are switched.
    """
    im1 = np.asarray(im1).astype(np.bool)
    im2 = np.asarray(im2).astype(np.bool)

    if im1.shape != im2.shape:
        raise ValueError("Shape mismatch: im1 and im2 must have the same shape.")

    im_sum = im1.sum() + im2.sum()
    if im_sum == 0:
        return empty_score

    # Compute Dice coefficient
    intersection = np.logical_and(im1, im2)

    return 2. * intersection.sum() / im_sum


def get_volume_dice_pair(mask_path, pred_path):
    pred = nib.load(pathlib.Path(pred_path).as_posix())
    mask = nib.load(pathlib.Path(mask_path).as_posix())

    pred = pred.get_fdata()
    mask = mask.get_fdata()

    volume = np.sum(mask)

    dice_value = dice(mask, pred)

    return volume, dice_value


def check_imgsize(json_path, tags, keys, indexes):
    min_x = 9999
    min_y = 9999
    min_z = 9999
    max_x = -1
    max_y = -1
    max_z = -1
    for item in yield_pathj(json_path, tags, keys, indexes):
        img = nib.load(item[0][0])
        zoom = img.get_fdata().shape
        x, y, z = zoom[0], zoom[1], zoom[2]

        min_x = min(min_x, x)
        min_y = min(min_y, y)
        min_z = min(min_z, z)

        max_x = max(max_x, x)
        max_y = max(max_y, y)
        max_z = max(max_z, z)
    return min_x, min_y


def check_imgsize_statistic(convert_json_path, tags, keys, input_index, output_index=-1):
    max_x_size = -1
    max_x_size_sample = None

    min_x_size = 9999
    min_x_size_sample = None

    max_y_size = -1
    max_y_size_sample = None

    min_y_size = 9999
    min_y_size_sample = None

    max_z_size = -1
    max_z_size_sample = None

    min_z_size = 9999
    min_z_size_sample = None

    min_volume = 9999
    min_volume_sample = None

    max_volume = -1
    max_volume_sample = None

    x_list = []
    y_list = []
    z_list = []
    vol_list = []

    for orig_paths, target_paths in _iter_convert_json(convert_json_path, tags, keys, input_index, output_index):
        orig_img_path = orig_paths[0]
        print(orig_img_path)

        img = nib.load(orig_img_path)
        zoom = img.get_fdata().shape
        x, y, z = zoom[0], zoom[1], zoom[2]
        volume = x * y * z

        x_list.append(x)
        y_list.append(y)
        z_list.append(z)
        vol_list.append(volume)

        if min_x_size_sample is None or x < min_x_size:
            min_x_size = x
            min_x_size_sample = orig_img_path
        if max_x_size_sample is None or x > max_x_size:
            max_x_size = x
            max_x_size_sample = orig_img_path
        if min_y_size_sample is None or y < min_y_size:
            min_y_size = y
            min_y_size_sample = orig_img_path
        if max_y_size_sample is None or y > max_y_size:
            max_y_size = y
            max_y_size_sample = orig_img_path
        if min_z_size_sample is None or z < min_z_size:
            min_z_size = z
            min_z_size_sample = orig_img_path
        if max_z_size_sample is None or z > max_z_size:
            max_z_size = z
            max_z_size_sample = orig_img_path
        if min_volume_sample is None or volume < min_volume:
            min_volume = volume
            min_volume_sample = orig_img_path
        if max_volume_sample is None or volume > max_volume:
            max_volume = volume
            max_volume_sample = orig_img_path

    print(f'max x: {max_x_size} at {max_x_size_sample}')
    print(f'max y: {max_y_size} at {max_y_size_sample}')
    print(f'max z: {max_z_size} at {max_z_size_sample}')
    print(f'max vol: {max_volume} at {max_volume_sample}')
    print(f'min x: {min_x_size} at {min_x_size_sample}')
    print(f'min y: {min_y_size} at {min_y_size_sample}')
    print(f'min z: {min_z_size} at {min_z_size_sample}')
    print(f'min vol: {min_volume} at {min_volume_sample}')
    print(f'avg x: {sum(x_list) / len(x_list)}')
    print(f'avg y: {sum(y_list) / len(y_list)}')
    print(f'avg z: {sum(z_list) / len(z_list)}')
    print(f'avg vol: {sum(vol_list) / len(vol_list)}')


def generate_new_unetr_dataset_json_from_convert_json(convert_json_path, new_json_save_path, keep_index):
    dataset = json.load(open(convert_json_path, 'r'))
    for tag in dataset.keys():
        for i in range(len(dataset[tag])):
            for key in dataset[tag][i].keys():
                dataset[tag][i][key] = dataset[tag][i][key][keep_index]

    json_file = json.dumps(dataset, indent=4, sort_keys=False)

    with open(new_json_save_path, 'w') as f:
        f.write(json_file)


def embed(small_array, big_array, big_index=None):
    """Overwrites values in big_array starting at big_index with those in small_array"""
    if big_index:
        slices = [np.s_[i:i+j] for i,j in zip(big_index, small_array.shape)]
        big_array[tuple(slices)] = small_array
        return big_array

def cover_rate(small_array_path, big_array_path):
    small_array = nib.load(small_array_path).get_fdata()
    small_voxel = np.sum(small_array)
    big_array = nib.load(big_array_path).get_fdata()
    big_voxel = np.sum(big_array)
    return small_voxel / big_voxel

# item['image'] = [img_0, img_1]
# item['label'] = [msk_0, msk_1]

# img_1 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_1 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# img_2 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_2 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# img_3 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'img' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
# msk_3 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'mask' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# #
# if tag == 'training':
#     img_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'img_train' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#     mak_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'mask_train' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# elif tag == 'validation':
#     img_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'img_val' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#     msk_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'mask_val' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# elif tag == 'test':
#     img_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'img_test' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#     msk_4 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_img_mask' + os.sep + 'mask_test' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# #
# img_5 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_5 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# if tag != 'test':
#     num = int(img.name.split('_')[0])
#     idx = cls_csv[cls_csv['MRN '] == num].index
#     try:
#         label = int(cls_csv.loc[idx, cls_col])
#         # wangqifeng_spacial_scale_intensity_cropped_96_224_224_classification
#         img_6 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm_classification' + os.sep + tag + os.sep + 'img' + os.sep + str(label) + '-' + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#         msk_6 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm_classification' + os.sep + tag + os.sep + 'mask' + os.sep + str(label) + '-' + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '.nii.gz').as_posix()
#     except:
#         img_6 = None
#         msk_6 = None
# else:
#     img_6 = None
#     msk_6 = None
# #
# img_7 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm_img_mask' + os.sep + 'img' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
# msk_7 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_scale_0-1500_without_norm_img_mask' + os.sep + 'mask' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# #
# if tag == 'training':
#     img_8 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'img_train' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#     msk_8 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'mask_train' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# elif tag == 'validation':
#     img_8 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'img_val' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#     msk_8 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'mask_val' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# else:
#     img_8 = None
#     msk_8 = None
# #
# if tag != 'test':
#     num = int(img.name.split('_')[0])
#     idx = cls_csv[cls_csv['MRN '] == num].index
#     try:
#         label = int(cls_csv.loc[idx, cls_col])
#         # wangqifeng_spacial_scale_intensity_cropped_96_224_224_classification
#         img_9 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_classification' + os.sep + tag + os.sep + 'img' + os.sep + str(label) + '-' + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#         msk_9 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_classification' + os.sep + tag + os.sep + 'mask' + os.sep + str(label) + '-' + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '.nii.gz').as_posix()
#     except:
#         img_9 = None
#         msk_9 = None
# else:
#     img_9 = None
#     mask_9 = None
# #
# img_10 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_10 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# img_11 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_scale_0-1500_without_norm' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_11 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_scale_0-1500_without_norm' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# img_12 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_12 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# img_13 = pathlib.Path(common_root.as_posix() + os.sep + '_temp' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine' + os.sep + img.parents[2].name + os.sep + img.parents[1].name + os.sep + img.parents[0].name + os.sep + img.name).as_posix()
# msk_13 = pathlib.Path(common_root.as_posix() + os.sep + '_temp' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine' + os.sep + msk.parents[2].name + os.sep + msk.parents[1].name + os.sep + msk.parents[0].name + os.sep + msk.name).as_posix()
# #
# if tag != 'test':
#     num = int(img.name.split('_')[0])
#     idx = cls_csv[cls_csv['MRN '] == num].index
#     try:
#         label = int(cls_csv.loc[idx, cls_col])
#         # wangqifeng_spacial_scale_intensity_cropped_96_224_224_classification
#         img_14 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_scale_0-1500_without_norm_classification' + os.sep + tag + os.sep + 'img' + os.sep + str(label) + '-' + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
#         msk_14 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_scale_0-1500_without_norm_classification' + os.sep + tag + os.sep + 'mask' + os.sep + str(label) + '-' + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '.nii.gz').as_posix()
#     except:
#         img_14 = None
#         msk_14 = None
# else:
#     img_14 = None
#     msk_14 = None
# #
# img_15 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'img' + os.sep + img.parents[2].name + '_' + img.parents[1].name + '_' + img.parents[0].name + '.nii.gz').as_posix()
# msk_15 = pathlib.Path(common_root.as_posix() + os.sep + '_out' + os.sep + 'wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask' + os.sep + 'mask' + os.sep + msk.parents[2].name + '_' + msk.parents[1].name + '_' + msk.parents[0].name + '_MASK.nii.gz').as_posix()
# #
# img_17 = None
# msk_17 = None
# item['image'] = [img_0, img_1, img_2, img_3, img_4, img_5, img_6, img_7, img_8, img_9, img_10, img_11, img_12, img_13, img_14, img_15, img_16]
# item['label'] = [msk_0, msk_1, msk_2, msk_3, msk_4, msk_5, msk_6, msk_7, msk_8, msk_9, msk_10, msk_11, msk_12, msk_13, msk_14, msk_15, msk_16]
